package com.feidee.fd.sml.algorithm.component.ml.regression

import com.feidee.fd.sml.algorithm.component.ml.MLParam
import org.apache.spark.ml.PipelineStage
import org.apache.spark.ml.regression.IsotonicRegression

/**
  * @Author: dongguosheng, songhaicheng
  * @Date: 2019/3/20 13:45
  * @Review songhaicheng
  */
case class IsotonicRegressionParam (
                                     override val input_pt: String,
                                     override val output_pt: String,
                                     override val hive_table: String,
                                     override val flow_time: String,
                                     override val featuresCol: String,
                                     override var labelCol: String,
                                     override var predictionCol: String,
                                     override val modelPath: String,
                                     override val metrics: Array[String],
                                     // 样本权重列
                                     weightCol: String,
                                     // 是否按升序输出保序结果，默认值 true
                                     isotonic: Boolean,
                                     // 特征数值所在的特征列（Vector 类型）索引值，如果特征列为 Double 类型，则此值不用设置，默认 0
                                     featureIndex: Int
                                   ) extends MLParam {
  def this() = this(null, null, null, null, "features", "label", "prediction", null, new Array[String](0),
    null, true, 0)

  override def verify(): Unit = {
    super.verify()
    require(featureIndex >= 0, "param featureIndex can't be negative")
  }

  override def toMap: Map[String, Any] = {
    var map = super.toMap
    map += ("weightCol" -> weightCol)
    map += ("isotonic" -> isotonic)
    map += ("featureIndex" -> featureIndex)
    map
  }
}


class IsotonicRegressionComponent extends AbstractRegressionComponent[IsotonicRegressionParam] {
  override def setUp(param: IsotonicRegressionParam): PipelineStage = {
    val isr = new IsotonicRegression()
      .setFeaturesCol(param.featuresCol)
      .setFeatureIndex(param.featureIndex)
      .setIsotonic(param.isotonic)
      .setLabelCol(param.labelCol)
      .setPredictionCol(param.predictionCol)

    if (tool.isNotNull(param.weightCol)) {
      isr.setWeightCol(param.weightCol)
    }

    isr
  }

}

object IsotonicRegressionComponent {
  def apply(paramStr: String): Unit = {
    new IsotonicRegressionComponent()(paramStr)
  }

  def main(args: Array[String]): Unit = {
    IsotonicRegressionComponent(args(0))
  }
}

